import tweb
import asyncio


class SessionBridge:
    def __init__(self, websocket_session):
        self.websocket_session = websocket_session
        self.tcp_reader = None
        self.tcp_writer = None
        self.tcp_to_ws_task = None

    async def connect_to_tcp_server(self, target_host, target_port):
        try:
            print(f"connecting target {target_host}:{target_port}")
            self.tcp_reader, self.tcp_writer = await asyncio.open_connection(target_host, int(target_port))
            self.websocket_session.send_text(f"[info] conneted to ssh target {target_host}:{target_port}\r\n")
            # 启动从TCP到WebSocket的转发
            self.tcp_to_ws_task = asyncio.create_task(self.forward_tcp_to_websocket())
            print(f"[info] begin forward ws <-> {target_host}:{target_port}")
            return True
        except Exception as e:
            print(f"连接到TCP服务器失败: {e}")
            if self.websocket_session:
                self.websocket_session.send_text(f"连接到target失败: {e}")
                self.websocket_session.close()
            return False

    async def forward_tcp_to_websocket(self):
        try:
            while True:
                data = await self.tcp_reader.read(4096)
                if not data:
                    print(f"[-] TCP连接已关闭")
                    break
                if self.websocket_session and not self.websocket_session.is_closed:
                    self.websocket_session.send_binary(data)
                else:
                    break
        except Exception as e:
            print(f"从TCP转发到WebSocket时出错: {e}")
        finally:
            await self.close()
            # 当TCP连接关闭时，也关闭WebSocket连接
            if self.websocket_session and not self.websocket_session.is_closed:
                print("[-] 关闭websocket连接")
                self.websocket_session.send_text("close websocket connection")
                self.websocket_session.close()

    async def forward_ws_to_tcp(self, message):
        try:
            if self.tcp_writer:
                if isinstance(message, str):
                    self.tcp_writer.write(message.encode('utf-8'))
                else:
                    self.tcp_writer.write(message)
                await self.tcp_writer.drain()
        except Exception as e:
            print(f"从WebSocket转发到TCP时出错: {e}")
            await self.close()

    async def close(self):
        # 关闭TCP连接
        if self.tcp_writer:
            self.tcp_writer.close()
            await self.tcp_writer.wait_closed()
            self.tcp_writer = None
        if self.tcp_to_ws_task:
            self.tcp_to_ws_task.cancel()
            self.tcp_to_ws_task = None
        if self.websocket_session and not self.websocket_session.is_closed:
            print(f"[-] 断开websocket连接: {self.websocket_session.id}")
            self.websocket_session.close('tcp connection closed')
            self.websocket_session = None


@tweb.websocket_handler("/ws/shell")
class WebShellHandler(tweb.WebsocketHandler):
    def __init__(self):
        # 保存所有活动的连接桥
        self.bridges = {}

    def on_open(self, session: tweb.WebsocketSession):
        # 为每个WebSocket会话创建一个独立的桥接
        print(f"[+] New WebSocket connection: {session.id}")
        print(f'request path: {session.request.parameter}')
        target_host, target_port = session.request.parameter.get("target", ":").split(":")
        if not target_host or not target_port:
            session.send_text("invalid target: host:port required")
            session.close()
            return
        bridge = SessionBridge(session)
        self.bridges[session.id] = bridge
        # 在异步环境中启动TCP连接
        asyncio.create_task(bridge.connect_to_tcp_server(target_host, target_port))

    def on_text_message(self, session, message):
        # 将文本消息转发到对应的TCP连接
        if session.id in self.bridges:
            bridge = self.bridges[session.id]
            asyncio.create_task(bridge.forward_ws_to_tcp(message))

    def on_binary_message(self, session, message):
        # 将二进制消息转发到对应的TCP连接
        if session.id in self.bridges:
            bridge = self.bridges[session.id]
            asyncio.create_task(bridge.forward_ws_to_tcp(message))

    def on_close(self, session, reason):
        # 关闭并清理对应的TCP连接
        print(f"[-] WebSocket connection closed: {session.id}")
        if session.id in self.bridges:
            bridge = self.bridges[session.id]
            asyncio.create_task(bridge.close())
            del self.bridges[session.id]